package cn.bugstack.hj.dev.tech.test;

import com.alibaba.fastjson.JSON;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@RunWith(SpringRunner.class)
@SpringBootTest
public class RAGTest {

    @Resource
    private OllamaChatClient ollamaChatClient;
    @Resource
    private TokenTextSplitter tokenTextSplitter;
    @Resource
    private SimpleVectorStore simpleVectorStore;
    @Resource
    private PgVectorStore pgVectorStore;

    @Test
    public void upload() {
        // ===== 1. 初始化文档阅读器 =====
        // 使用Apache Tika库创建文档阅读器，支持解析多种格式（PDF/Word/TXT等）
        // 注意："./data/file.txt"是相对路径，建议改为"classpath:data/file.txt"确保测试可移植性
        TikaDocumentReader reader = new TikaDocumentReader("data/file.txt");
        // ===== 2. 读取并解析文档 =====
        // 调用Tika解析器提取文档内容，返回Document对象列表
        // 每个Document包含：
        //   - content: 文本内容
        //   - metadata: 元数据（如文件名、格式、创建时间等）
        List<Document> documents = reader.get();
        // ===== 3. 文本分块处理 =====
        // 使用预配置的tokenTextSplitter对文档进行分块：
        //   - 按固定token数分割（如512）
        //   - 保留段落完整性
        //   - 添加重叠文本防止信息断裂
        // 目的：适配大模型输入长度限制
        List<Document> documentSplitterList = tokenTextSplitter.apply(documents);
        // ===== 4. 添加知识库元数据（原始文档） =====
        // 为原始文档添加"knowledge"元数据字段，标记所属知识库名称
        // 用途：后续检索时区分不同知识库来源
        documents.forEach(doc -> doc.getMetadata().put("knowledge", "知识库名称"));
        // ===== 5. 添加知识库元数据（分块后文档） =====
        // 同样为分块后的文档添加知识库标记，确保信息可溯源
        documentSplitterList.forEach(doc -> doc.getMetadata().put("knowledge", "知识库名称"));
        // ===== 6. 存储到向量数据库 =====
        // 将处理后的文档块存入PostgreSQL向量数据库：
        //   1. 通过嵌入模型（如OpenAI Embeddings）将文本转为向量
        //   2. 向量和元数据一并存入PGVector
        //   3. 自动创建/更新向量索引
        pgVectorStore.accept(documentSplitterList);
        // ===== 7. 日志记录 =====
        // 打印操作完成日志（实际项目建议添加更详细的调试信息）
        log.info("上传完成");
    }


    @Test
    public void chat() {
        // ===== 1. 用户问题输入 =====
        String message = "黄哈哈，哪年出生";
        // ===== 2. 系统提示词模板 =====
        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;
        // ===== 3. 构建向量搜索请求 =====
        // 创建搜索请求对象，配置：
        // - 查询文本：用户问题
        // - 返回结果数：5条（withTopK）
        // - 过滤条件：只检索"知识库名称"标签的内容（withFilterExpression）
        SearchRequest request = SearchRequest.query(message).withTopK(5).withFilterExpression("knowledge == '知识库名称'");
        // ===== 4. 执行向量相似度搜索 =====
        // 从PGVector向量库中检索相关文档片段
        // 返回结果按与问题的语义相似度排序
        List<Document> documents = pgVectorStore.similaritySearch(request);
        // ===== 5. 合并检索结果 =====
        // 将多个文档片段内容合并为单个字符串，用空字符串连接
        // 用于后续插入系统提示模板
        String documentsCollectors = documents.stream().map(Document::getContent).collect(Collectors.joining());
        // ===== 6. 构建系统提示消息 =====
        // 将检索到的文档内容注入系统提示模板
        // {documents}占位符会被实际内容替换
        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentsCollectors));
        // ===== 7. 组装对话消息 =====
        // 按时间顺序组织消息列表：
        // 1. 用户原始问题
        // 2. 包含知识库内容的系统提示
        ArrayList<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(message));
        messages.add(ragMessage);
        // ===== 8. 调用大模型生成回答 =====
        // 使用Ollama接口调用本地模型（deepseek-r1:1.5b）：
        // - 输入组装好的消息列表
        // - 返回包含模型生成的响应
        ChatResponse chatResponse = ollamaChatClient.call(new Prompt(messages, OllamaOptions.create().withModel("deepseek-r1:1.5b")));
        // ===== 9. 输出结果 =====
        // 将响应对象转为JSON格式打印（便于查看完整结构）
        log.info("测试结果:{}", JSON.toJSONString(chatResponse));

    }

}
